/*
 * FILE: rdt_receiver.cc
 * DESCRIPTION: Reliable data transfer receiver.
 * NOTE: This implementation assumes there is no packet loss, corruption, or 
 *       reordering.  You will need to enhance it to deal with all these 
 *       situations.  In this implementation, the packet format is laid out as 
 *       the following:
 *       
 *       |<-  1 byte  ->|<-             the rest            ->|
 *       | payload size |<-             payload             ->|
 *
 *       The first byte of each packet indicates the size of the payload
 *       (excluding this single-byte header)
 */


#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <map>
#include <iostream>

#include "rdt_struct.h"
#include "rdt_receiver.h"

std::map<int, std::string> receiver_saved;
int max_received_index = 0;
int done_index = 0; // Index start from 1

void send_acknowledge(int index);

/* receiver initialization, called once at the very beginning */
void Receiver_Init() {
    fprintf(stdout, "At %.2fs: receiver initializing ...\n", GetSimulationTime());
}

/* receiver finalization, called once at the very end.
   you may find that you don't need it, in which case you can leave it blank.
   in certain cases, you might want to use this opportunity to release some 
   memory you allocated in Receiver_init(). */
void Receiver_Final() {
    fprintf(stdout, "At %.2fs: receiver finalizing ...\n", GetSimulationTime());
}

/* event handler, called when a packet is passed from the lower layer at the 
   receiver */
void Receiver_FromLowerLayer(struct packet *pkt) {
#if USE_ORG
    /* 1-byte header indicating the size of the payload */
    int header_size = 1;

    /* construct a message and deliver to the upper layer */
    struct message *msg = (struct message *) malloc(sizeof(struct message));
    ASSERT(msg != NULL);

    msg->size = pkt->data[0];

    /* sanity check in case the packet is corrupted */
    if (msg->size < 0) msg->size = 0;
    if (msg->size > RDT_PKTSIZE - header_size) msg->size = RDT_PKTSIZE - header_size;

    msg->data = (char *) malloc(msg->size);
    ASSERT(msg->data != NULL);
    memcpy(msg->data, pkt->data + header_size, msg->size);
    Receiver_ToUpperLayer(msg);

    /* don't forget to free the space */
    if (msg->data != NULL) free(msg->data);
    if (msg != NULL) free(msg);
#else
    int payload_size = pkt->data[0];
    int real_size = payload_size / 2 - INDEX_SIZE;

    const int max_payload_size = RDT_PKTSIZE - HEADER_SIZE;
    if (payload_size <= 0 || payload_size > max_payload_size) {
        return;
    }

    std::string str(pkt->data + HEADER_SIZE, payload_size);
    int index;

    std::string str_msg = DecodeMessage(str, index);
    if (str_msg.empty()) {
        // Corrupted
        // std::cout << "Packet damaged." << std::endl;
        return;
    }

    // Has been received
    if (index <= done_index || receiver_saved.find(index) != receiver_saved.end()) {
        send_acknowledge(index);
        return;
    }

    // Add to buffer
    receiver_saved.insert(std::make_pair(index, str_msg.substr(0, real_size)));

    // std::cout << "R " << real_size << " " << str_msg << std::endl;
    // std::cout << "Receive index " << index << std::endl;

    if (index > max_received_index) {
        max_received_index = index;
    }
    if (index == done_index + 1) {
        std::map<int, std::string>::const_iterator it;

        // Transfer continuous index
        int cur_index = index;
        while ((it = receiver_saved.find(cur_index)) != receiver_saved.end()) {
            std::string cur_msg = it->second;

            // New Message
            message msg{};
            msg.size = cur_msg.size();
            msg.data = new char[msg.size];
            memcpy(msg.data, cur_msg.data(), msg.size);

            // std::cout << "Transfer to upper " << cur_index << std::endl;
            Receiver_ToUpperLayer(&msg);
            send_acknowledge(cur_index);

            // Delete Message
            delete[]msg.data;

            done_index = cur_index;
            cur_index++;
            receiver_saved.erase(it);
        }
    }

#endif
}

void send_acknowledge(int index) {
    std::string data(MESSAGE_ACK);

    int max_payload_size = RDT_PKTSIZE - HEADER_SIZE;
    int data_size = max_payload_size / 2 - INDEX_SIZE;

    int real_size = data.size();
    if (real_size < data_size) {
        data.resize(data_size, 0);
    }
    int payload_size = (real_size + INDEX_SIZE) * 2;

    std::string encode_res = EncodeMessage(data, index);

    // Copy data to packet
    packet pkt{};
    memset(pkt.data, 0, RDT_PKTSIZE);

    pkt.data[0] = char(payload_size);
    memcpy(pkt.data + HEADER_SIZE, encode_res.data(), payload_size);

    Receiver_ToLowerLayer(&pkt);

    // std::cout << "Receiver send ack " << index << std::endl;
}